// Maintained by Shenglin Qin, Chengfu Zou
// Copyright (C) FYT Vision Group. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "armor_detector/light_corner_corrector.hpp"

#include <numeric>

namespace fyt::auto_aim {

void LightCornerCorrector::correctCorners(Armor &armor, const cv::Mat &gray_img) {
  // If the width of the light is too small, the correction is not performed
  constexpr int PASS_OPTIMIZE_WIDTH = 3;

  if (armor.left_light.width > PASS_OPTIMIZE_WIDTH) {
    // Find the symmetry axis of the light
    SymmetryAxis left_axis = findSymmetryAxis(gray_img, armor.left_light);
    // Find the corner of the light
    armor.left_light.top = findCorner(gray_img, armor.left_light, left_axis, "top");
    armor.left_light.bottom = findCorner(gray_img, armor.left_light, left_axis, "bottom");
    armor.left_light.center = left_axis.centroid;
  }

  if (armor.right_light.width > PASS_OPTIMIZE_WIDTH) {
    // Find the symmetry axis of the light
    SymmetryAxis right_axis = findSymmetryAxis(gray_img, armor.right_light);
    // Find the corner of the light
    armor.right_light.top = findCorner(gray_img, armor.right_light, right_axis, "top");
    armor.right_light.bottom = findCorner(gray_img, armor.right_light, right_axis, "bottom");
    armor.right_light.center = right_axis.centroid;
  }
}


SymmetryAxis LightCornerCorrector::findSymmetryAxis(const cv::Mat &gray_img, const Light &light) {
  constexpr float MAX_BRIGHTNESS = 25;
  constexpr float SCALE = 0.07;

  // Scale the bounding box
  cv::Rect light_box = light.boundingRect();
  light_box.x -= light_box.width * SCALE;
  light_box.y -= light_box.height * SCALE;
  light_box.width += light_box.width * SCALE * 2;
  light_box.height += light_box.height * SCALE * 2;

  // Check boundary
  light_box.x = std::max(light_box.x, 0);
  light_box.x = std::min(light_box.x, gray_img.cols - 1);
  light_box.y = std::max(light_box.y, 0);
  light_box.y = std::min(light_box.y, gray_img.rows - 1);
  light_box.width = std::min(light_box.width, gray_img.cols - light_box.x);
  light_box.height = std::min(light_box.height, gray_img.rows - light_box.y);

  // Get normalized light image
  cv::Mat roi = gray_img(light_box);
  float mean_val = cv::mean(roi)[0];
  roi.convertTo(roi, CV_32F);
  cv::normalize(roi, roi, 0, MAX_BRIGHTNESS, cv::NORM_MINMAX);

  // Calculate the centroid
  cv::Moments moments = cv::moments(roi, false);
  cv::Point2f centroid = cv::Point2f(moments.m10 / moments.m00, moments.m01 / moments.m00) +
                         cv::Point2f(light_box.x, light_box.y);

  // Initialize the PointCloud
  std::vector<cv::Point2f> points;
  for (int i = 0; i < roi.rows; i++) {
    for (int j = 0; j < roi.cols; j++) {
      for (int k = 0; k < std::round(roi.at<float>(i, j)); k++) {
        points.emplace_back(cv::Point2f(j, i));
      }
    }
  }
  cv::Mat points_mat = cv::Mat(points).reshape(1);

  // PCA (Principal Component Analysis)
  auto pca = cv::PCA(points_mat, cv::Mat(), cv::PCA::DATA_AS_ROW);

  // Get the symmetry axis
  cv::Point2f axis =
    cv::Point2f(pca.eigenvectors.at<float>(0, 0), pca.eigenvectors.at<float>(0, 1));

  // Normalize the axis
  axis = axis / cv::norm(axis);

  return SymmetryAxis{.centroid = centroid, .direction = axis, .mean_val = mean_val};
}


/*
  SymmetryAxis LightCornerCorrector::findSymmetryAxis(
      const cv::Mat &gray_img, const Light &light) {

    constexpr float MAX_BRIGHTNESS = 25.0f;
    constexpr float SCALE          = 0.07f;

    //—— 1. 扩张并截取 ROI ————————————————————————————————
    cv::Rect light_box = light.boundingRect();
    light_box.x      -= light_box.width  * SCALE;
    light_box.y      -= light_box.height * SCALE;
    light_box.width  += light_box.width  * SCALE * 2;
    light_box.height += light_box.height * SCALE * 2;
    light_box &= cv::Rect(0, 0, gray_img.cols, gray_img.rows);

    cv::Mat roi = gray_img(light_box);

    //— 2. 计算原始均值（供外部阈值或判断使用） —————————————
    float mean_val = static_cast<float>(cv::mean(roi)[0]);

    //— 3. 归一化到 [0, MAX_BRIGHTNESS] 并转为 float ——————————
    cv::Mat f;
    roi.convertTo(f, CV_32F);
    cv::normalize(f, f, 0.0f, MAX_BRIGHTNESS, cv::NORM_MINMAX);

    //— 4. 用 cv::moments 计算矩（包括 m00, m10, m01, mu20, mu11, mu02）—
    cv::Moments mo = cv::moments(f,false);

    // 避免 m00=0 导致除零
    if (mo.m00 < 1e-6) {
      return SymmetryAxis{
        .centroid = light.center,
        .direction= cv::Point2f(1, 0),
        .mean_val = mean_val
      };
    }

    //—— 5. 质心（ROI 坐标系→全局坐标系）———————————————
    float cx = static_cast<float>(mo.m10 / mo.m00);
    float cy = static_cast<float>(mo.m01 / mo.m00);
    cv::Point2f centroid(cx + light_box.x, cy + light_box.y);

    //—— 6. 中心矩归一化后求特征方向 θ = ½ atan2(2μ11, μ20−μ02) ——————
    double mu20 = mo.mu20 / mo.m00;
    double mu11 = mo.mu11 / mo.m00;
    double mu02 = mo.mu02 / mo.m00;
    double theta = 0.5 * std::atan2(2.0 * mu11, mu20 - mu02);

    cv::Point2f axis(
      static_cast<float>(std::cos(theta)),
      static_cast<float>(std::sin(theta))
    );

    return SymmetryAxis{
      .centroid = centroid,
      .direction= axis,
      .mean_val = mean_val
    };
  }
*/




cv::Point2f LightCornerCorrector::findCorner(const cv::Mat &gray_img,
                                             const Light &light,
                                             const SymmetryAxis &axis,
                                             std::string order) {
  constexpr float START = 0.8 / 2;
  constexpr float END = 1.2 / 2;

  auto inImage = [&gray_img](const cv::Point &point) -> bool {
    return point.x >= 0 && point.x < gray_img.cols && point.y >= 0 && point.y < gray_img.rows;
  };

  auto distance = [](float x0, float y0, float x1, float y1) -> float {
    return std::sqrt((x0 - x1) * (x0 - x1) + (y0 - y1) * (y0 - y1));
  };

  int oper = order == "top" ? 1 : -1;
  float L = light.length;
  float dx = axis.direction.x * oper;
  float dy = axis.direction.y * oper;

  std::vector<cv::Point2f> candidates;
  
  // Select multiple corner candidates and take the average as the final corner
  int n = light.width - 2;
  int half_n = std::round(n / 2);
  for (int i = -half_n; i <= half_n; i++) {
    float x0 = axis.centroid.x + L * START * dx + i;
    float y0 = axis.centroid.y + L * START * dy;

    cv::Point2f prev = cv::Point2f(x0, y0);
    cv::Point2f corner = cv::Point2f(x0, y0);
    float max_brightness_diff = 0;
    // Search along the symmetry axis to find the corner that has the maximum brightness difference
    for (float x = x0 + dx, y = y0 + dy; distance(x, y, x0, y0) < L * (END - START);
         x += dx, y += dy) {
      cv::Point2f cur = cv::Point2f(x, y);
      if (!inImage(cv::Point(cur))) {
        break;
      }

      float brightness_diff = gray_img.at<uchar>(cv::Point(prev)) - gray_img.at<uchar>(cur);
      if (brightness_diff > max_brightness_diff && gray_img.at<uchar>(prev) > axis.mean_val) {
        max_brightness_diff = brightness_diff;
        corner = prev;
      }

      prev = cur;
    }

    candidates.emplace_back(corner);
  }
  cv::Point2f result = std::accumulate(candidates.begin(), candidates.end(), cv::Point2f(0, 0)) /
                       static_cast<float>(candidates.size());

  return result;
}

/*  
  cv::Point2f LightCornerCorrector::findCorner(
      const cv::Mat &gray_img,
      const Light &light,
      const SymmetryAxis &axis,
      std::string order) {
    constexpr float START = 0.8f / 2;
    constexpr float END   = 1.2f / 2;

    // 判断点是否在图像范围内
    auto inImage = [&gray_img](const cv::Point2f &pt) {
      return pt.x >= 0 && pt.x < gray_img.cols && pt.y >= 0 && pt.y < gray_img.rows;
    };

    int oper      = (order == "top") ? 1 : -1;
    float L       = light.length;
    float dx      = axis.direction.x * oper;
    float dy      = axis.direction.y * oper;

    // 候选角点容器
    std::vector<cv::Point2f> candidates;
    candidates.reserve(3);

    // 取中间 3 个扫描起点
    constexpr int CAND_CNT = 3;
    int half_n = CAND_CNT / 2;
    for (int i = -half_n; i <= half_n; ++i) {
      // 初始点
      cv::Point2f prev(
        axis.centroid.x + L * START * dx + static_cast<float>(i),
        axis.centroid.y + L * START * dy
      );
      cv::Point2f corner = prev;
      float max_diff     = 0.0f;

      // 搜索最大亮度差位置
      int max_steps = static_cast<int>(L * (END - START));
      cv::Point2f cur = prev;
      for (int step = 0; step < max_steps; ++step) {
        cur.x += dx;
        cur.y += dy;
        if (!inImage(cur)) break;

        // 整数坐标访问
        int px = static_cast<int>(std::round(prev.x));
        int py = static_cast<int>(std::round(prev.y));
        int cx = static_cast<int>(std::round(cur.x));
        int cy = static_cast<int>(std::round(cur.y));
        uchar prev_val = gray_img.ptr<uchar>(py)[px];
        uchar cur_val  = gray_img.ptr<uchar>(cy)[cx];

        float diff = static_cast<float>(prev_val) - static_cast<float>(cur_val);
        if (diff > max_diff && prev_val > axis.mean_val) {
          max_diff = diff;
          corner   = prev;
        }
        prev = cur;
      }
      candidates.push_back(corner);
    }

    // 平均候选点
    cv::Point2f sum(0.0f, 0.0f);
    for (auto &pt : candidates) sum += pt;
    return sum * (1.0f / candidates.size());
  }
  
*/


}  // namespace fyt::auto_aim
